#!/usr/bin/env python3
import os
import sys

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/')
)

import grpc
import struct
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from time import sleep
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections


eth_header_format = "!6s6sH"
cpu_header_format = "!B"

eth_header_length = struct.calcsize(eth_header_format)
cpu_header_offset = eth_header_length

CLONE_SESSION_ID = 100
CPU_PORT = 252

def set_ecn_threshold(p4info_helper, switch, threshold_value):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecn_config",
        match_fields={},
        action_name="MyIngress.set_ecn_threshold",
        action_params={"threshold": threshold_value}
    )
    switch.WriteTableEntry(table_entry)
    print(f"[INFO] ECN threshold set to {threshold_value} on {switch.name}")

def install_ipv4_lpm(p4info_helper, switch, dst_ip, dst_mac, out_port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip, 32)},
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": out_port
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"[{switch.name}] Rule: {dst_ip} -> {dst_mac} via port {out_port}")


def writeCheck_ecnRules(p4info_helper, sw, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.check_ecn",
        action_name="MyEgress.mark_ecn",
        default_action=True
    )
    sw.WriteTableEntry(table_entry)


def writeCloneRules(p4info_helper, sw, clone_session_id, replicas):
    session_entry = p4info_helper.buildCloneSessionEntry(
        clone_session_id=clone_session_id,
        replicas=replicas
    )
    sw.WritePREEntry(session_entry)


def fetch_responses(connection):
    try:
        for response in connection.stream_msg_resp:
            print("Received response - - - - -")
            if response.WhichOneof("update") == "packet":
                packet = response.packet.payload
                cpu_header = packet[cpu_header_offset:cpu_header_offset+1]
                ecn = struct.unpack(cpu_header_format, cpu_header)[0]
                print("ECN value:", ecn)
                if ecn == 3:
                    print("Congestion happens!!")
    except grpc.RpcError as e:
        pass


def main():
    p4info_path = "build/ecn.p4.p4info.txt"
    bmv2_json_path = "build/ecn.json"

    try:
        threshold = int(input("Please input the threshold of the queue: "))
    except ValueError:
        print("[ERROR] Please enter an integer for threshold.")
        return

    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_path)

    try:
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-ecn.txt'
        )
        s1.MasterArbitrationUpdate()
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
        set_ecn_threshold(p4info_helper, s1, threshold)

        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-ecn.txt'
        )
        s2.MasterArbitrationUpdate()
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
        set_ecn_threshold(p4info_helper, s2, threshold)

        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-ecn.txt'
        )
        s3.MasterArbitrationUpdate()
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
        set_ecn_threshold(p4info_helper, s3, threshold)

        # 安装转发表
        install_ipv4_lpm(p4info_helper, s1, "10.0.1.1", "08:00:00:00:01:01", 1)
        install_ipv4_lpm(p4info_helper, s1, "10.0.1.11", "08:00:00:00:01:11", 2)
        install_ipv4_lpm(p4info_helper, s1, "10.0.2.2", "08:00:00:00:02:02", 3)
        install_ipv4_lpm(p4info_helper, s1, "10.0.2.22", "08:00:00:00:02:22", 3)
        install_ipv4_lpm(p4info_helper, s1, "10.0.3.3", "08:00:00:00:03:03", 4)

        install_ipv4_lpm(p4info_helper, s2, "10.0.2.2", "08:00:00:00:02:02", 2)
        install_ipv4_lpm(p4info_helper, s2, "10.0.2.22", "08:00:00:00:02:22", 1)
        install_ipv4_lpm(p4info_helper, s2, "10.0.1.1", "08:00:00:00:01:01", 3)
        install_ipv4_lpm(p4info_helper, s2, "10.0.1.11", "08:00:00:00:01:11", 3)
        install_ipv4_lpm(p4info_helper, s2, "10.0.3.3", "08:00:00:00:03:03", 4)

        install_ipv4_lpm(p4info_helper, s3, "10.0.3.3", "08:00:00:00:03:03", 1)
        install_ipv4_lpm(p4info_helper, s3, "10.0.1.1", "08:00:00:00:01:01", 2)
        install_ipv4_lpm(p4info_helper, s3, "10.0.1.11", "08:00:00:00:01:11", 2)
        install_ipv4_lpm(p4info_helper, s3, "10.0.2.2", "08:00:00:00:02:02", 3)
        install_ipv4_lpm(p4info_helper, s3, "10.0.2.22", "08:00:00:00:02:22", 3)

        # 安装ECN检测和Clone规则
        writeCheck_ecnRules(p4info_helper, s1, {"ecn_threshold": threshold})
        writeCheck_ecnRules(p4info_helper, s2, {"ecn_threshold": threshold})
        writeCheck_ecnRules(p4info_helper, s3, {"ecn_threshold": threshold})

        writeCloneRules(p4info_helper, s1, 100, [{"egress_port": 252, "instance": 1}])
        writeCloneRules(p4info_helper, s2, 100, [{"egress_port": 252, "instance": 1}])
        writeCloneRules(p4info_helper, s3, 100, [{"egress_port": 252, "instance": 1}])

        print("\nMonitoring network congestion...")

        while True:
            fetch_responses(s1)
            fetch_responses(s2)
            fetch_responses(s3)

    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()

if __name__ == '__main__':
    main()
